from pymilvus import (
    connections,
    FieldSchema,
    CollectionSchema,
    DataType,
    Collection,
    utility
)
from typing import List, Dict
import numpy as np


def write_to_milvus(structured_data: List[Dict], embeddings: List[List[float]],
                    collection_name: str = "labor_contract_law"):
    """
    将法律文本结构化数据及向量写入本地Milvus数据库（修复TEXT类型兼容性问题）

    参数:
        structured_data: 结构化元数据列表（上一步骤输出）
        embeddings: 向量列表（text_embeddings函数输出）
        collection_name: Milvus集合名称，默认"labor_contract_law"
    """
    # 1. 连接本地Milvus服务（默认端口）
    try:
        connections.connect(
            alias="default",
            host="localhost",  # 本地部署地址
            port="19530"  # 默认端口
        )
        print("成功连接到本地Milvus服务")
    except Exception as e:
        raise ConnectionError(f"连接Milvus失败：{str(e)}")

    # 2. 验证输入数据完整性
    if len(structured_data) != len(embeddings):
        raise ValueError(f"结构化数据与向量数量不匹配：{len(structured_data)} vs {len(embeddings)}")

    # 验证向量维度一致性
    vector_dim = len(embeddings[0]) if embeddings else 0
    for i, emb in enumerate(embeddings):
        if len(emb) != vector_dim:
            raise ValueError(f"向量{i}维度异常：实际{len(emb)}维，预期{vector_dim}维")

    # 3. 定义集合结构（适配法律文本元数据，使用VARCHAR替代TEXT以兼容旧版本）
    fields = [
        # 主键：法条唯一ID（如labor_law_001）
        FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=50, is_primary=True),
        # 向量字段：1024维向量
        FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=vector_dim),
        # 法律名称（如"中华人民共和国劳动合同法"）
        FieldSchema(name="law_name", dtype=DataType.VARCHAR, max_length=100),
        # 章节名称（如"第二章 劳动合同的订立"）
        FieldSchema(name="chapter", dtype=DataType.VARCHAR, max_length=100),
        # 章节编码（如"02"）
        FieldSchema(name="chapter_code", dtype=DataType.VARCHAR, max_length=10),
        # 法条编号（如"第十条"）
        FieldSchema(name="article_num", dtype=DataType.VARCHAR, max_length=20),
        # 法条编码（如"010"）
        FieldSchema(name="article_num_code", dtype=DataType.VARCHAR, max_length=10),
        # 法条标题
        FieldSchema(name="article_title", dtype=DataType.VARCHAR, max_length=200),
        # 清理后的法条内容（使用VARCHAR替代TEXT，设置较大长度）
        FieldSchema(name="content_clean", dtype=DataType.VARCHAR, max_length=4096),  # 关键修改点
        # 适用场景标签（逗号分隔字符串）
        FieldSchema(name="applicable_scenarios", dtype=DataType.VARCHAR, max_length=500),
        # 检索权重
        FieldSchema(name="retrieval_weight", dtype=DataType.FLOAT)
    ]

    # 创建集合 schema
    schema = CollectionSchema(
        fields=fields,
        description="中华人民共和国劳动合同法法条向量集合",
        enable_dynamic_field=False  # 关闭动态字段，确保数据结构一致
    )

    # 4. 创建或加载集合
    if utility.has_collection(collection_name, using="default"):
        # 若集合已存在，先删除旧集合（因为字段类型有变化）
        utility.drop_collection(collection_name, using="default")
        print(f"已删除旧集合：{collection_name}")

    # 创建新集合
    collection = Collection(name=collection_name, schema=schema, using="default")
    print(f"已创建新集合：{collection_name}")

    # 创建向量索引（优化检索性能）
    index_params = {
        "index_type": "IVF_FLAT",  # 适合中小规模数据集（<100万）
        "metric_type": "L2",  # 欧氏距离（法律文本语义匹配常用）
        "params": {"nlist": 128}  # 聚类数量，根据数据量调整
    }
    collection.create_index(
        field_name="embedding",
        index_params=index_params
    )
    print("已创建向量索引")

    # 5. 准备插入数据
    insert_data = [
        # 提取结构化数据中的字段，与schema对应
        [item["id"] for item in structured_data],
        embeddings,  # 向量列表
        [item["metadata"]["law_name"] for item in structured_data],
        [item["metadata"]["chapter"] for item in structured_data],
        [item["metadata"]["chapter_code"] for item in structured_data],
        [item["metadata"]["article_num"] for item in structured_data],
        [item["metadata"]["article_num_code"] for item in structured_data],
        [item["metadata"]["article_title"] for item in structured_data],
        [item["metadata"]["content_clean"] for item in structured_data],
        [",".join(item["metadata"]["applicable_scenarios"]) for item in structured_data],  # 列表转字符串
        [item["metadata"]["retrieval_weight"] for item in structured_data]
    ]

    # 6. 批量插入数据
    try:
        result = collection.insert(insert_data)
        # 刷新集合使数据可查
        collection.flush()
        print(f"成功插入{len(result.primary_keys)}条数据，主键列表：{result.primary_keys[:5]}...")
    except Exception as e:
        raise RuntimeError(f"数据插入失败：{str(e)}")

    # 7. 关闭连接
    connections.disconnect("default")
    return True


# 测试函数
if __name__ == "__main__":
    # 模拟结构化数据（上一步输出）
    sample_structured = [
        {
            "id": "labor_law_010",
            "vector_text": "第十条 建立劳动关系，应当订立书面劳动合同...",
            "metadata": {
                "law_name": "中华人民共和国劳动合同法",
                "chapter": "第二章 劳动合同的订立",
                "chapter_code": "02",
                "article_num": "第十条",
                "article_num_code": "010",
                "article_title": "书面劳动合同订立要求",
                "content_clean": "建立劳动关系，应当订立书面劳动合同。已建立劳动关系，未同时订立书面劳动合同的，应当自用工之日起一个月内订立书面劳动合同。用人单位与劳动者在用工前订立劳动合同的，劳动关系自用工之日起建立。",
                "applicable_scenarios": ["contract_signing", "employment"],
                "retrieval_weight": 1.2
            }
        }
    ]

    # 模拟向量（1024维）
    sample_embeddings = [np.random.rand(1024).tolist()]

    try:
        write_to_milvus(sample_structured, sample_embeddings)
        print("测试写入成功")
    except Exception as e:
        print(f"测试失败：{e}")
